//
// Created by hjie on 23-7-4.
//
#include <string>
#include <cstring>
#include <unistd.h>

extern "C"
{
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libavutil/error.h>
#define __STDC_FORMAT_MACROS

}

void save_gray_frame(unsigned char * buffer, int wrap, int xsize, int ysize, const char * filename)
{
    FILE * fp = NULL;
    fp = fopen(filename, "w");
    if (fp == NULL)
    {
        printf("[ERROR][fopen] failed filename : %s.\n", filename);
        return;
    }
    fprintf(fp, "P5\n%d %d\n%d\n", xsize, ysize, 255);
    for (int index = 0; index < ysize; index++)
    {
        fwrite(buffer + index * wrap ,1, xsize, fp);
    }
    fclose(fp);
}

void save_yuv_frame(
        unsigned char * buffer_0,
        unsigned char * buffer_1,
        unsigned char * buffer_2,
        int wrap_0,
        int wrap_1,
        int wrap_2, int xsize, int ysize, const char * filename)
{
    FILE * fp = NULL;
    fp = fopen(filename, "w");
    if (fp == NULL)
    {
        printf("[ERROR][fopen] failed filename : %s.\n", filename);
        return;
    }
    int width  = xsize/2;
    int height = ysize/2;
    fprintf(fp, "P5\n%d %d\n%d\n", width * 2, height * 2, 255);
    for (int index = 0; index < height * 2; index++)
    {
        fwrite(buffer_0 + index * wrap_0, width * 2, 1, fp);
    }
    for (int index = 0; index < height; index++)
    {
        fwrite(buffer_1 + index * wrap_1, width, 1, fp);
    }
    for (int index = 0; index < height; index++)
    {
        fwrite(buffer_2 + index * wrap_2, width, 1, fp);
    }
}

int decode_packet(AVPacket * packet, AVCodecContext * codec_context, AVFrame * frame)
{
    int response = avcodec_send_packet(codec_context, packet);
    if (response < 0)
    {
        printf("[ERROR][avcodec_send_packet] while sending a packet to the decoder: %d.\n", response);
        return response;
    }
    while(response >= 0)
    {
        response = avcodec_receive_frame(codec_context, frame);
        if (response == AVERROR(EAGAIN) || response == AVERROR_EOF)
        {
            break;
        }
        else if (response < 0)
        {
            printf("[ERROR][avcodec_receive_frame] while receiving a frame from the decoder: %d.\n", response);
            return response;
        }
        if (response >= 0)
        {
            printf("Frame %d (type=%c, size=%d bytes, format %d) pts %ld key_frame %d [DTS %d].\n",
                   codec_context->frame_number,
                   av_get_picture_type_char(frame->pict_type),
                   frame->pkt_size,
                   frame->format,
                   frame->pts,
                   frame->key_frame,
                   frame->coded_picture_number);
        }
        char frame_filename[128] = {0};
        // portable gray map
        snprintf(frame_filename, sizeof(frame_filename) - 1, "%s-%d.pgm", "frame", codec_context->frame_number);
        if (frame->format != AV_PIX_FMT_YUV420P)
        {
            printf("[Warning] the generated file may not be a grayscale image, but could e.g. be just R component if the video format is RGB.\n");
        }
        save_gray_frame(frame->data[0], frame->linesize[0], frame->width, frame->height, frame_filename);
        memset(frame_filename, 0, sizeof(frame_filename));
        snprintf(frame_filename, sizeof(frame_filename) - 1, "%s-uv-%d.pgm", "frame", codec_context->frame_number);
        save_yuv_frame(frame->data[0],
                       frame->data[1],
                       frame->data[2],
                       frame->linesize[0],
                       frame->linesize[1],
                       frame->linesize[2],
                       frame->width, frame->height, frame_filename);
    }
    return 0;
}

int main(int argc, const char * argv[])
{
    if (argc < 2)
    {
        printf("You need to specify a media file.\n");
        return -1;
    }
    std::string filename = argv[1];
    AVFormatContext * format_context = avformat_alloc_context();
    int result = 0;
    result = avformat_open_input(&format_context, filename.c_str(), NULL, NULL);
    if (result < 0)
    {
        printf("[ERROR][avformat_open_input] result : %d.\n", result);
        return -1;
    }
    result = avformat_find_stream_info(format_context, NULL);
    if (result < 0)
    {
        printf("[ERROR][avformat_find_stream_info] could not get stream information result : %d.\n", result);
        return -1;
    }
    printf("Format : %s, duration %ld us", format_context->iformat->long_name, format_context->duration);

    const AVCodec * codec = NULL;
    AVCodecParameters * codec_parameters = NULL;
    int video_stream_index = -1;
    for (unsigned int index = 0; index < format_context->nb_streams; index++)
    {
        AVCodecParameters * local_parameters = NULL;
        local_parameters = format_context->streams[index]->codecpar;

        printf("AVStream->time_base before open coded %d/%d.\n",
               format_context->streams[index]->time_base.num,format_context->streams[index]->time_base.den);
        printf("AVStream->r_frame_rate before open coded %d/%d.\n",
               format_context->streams[index]->r_frame_rate.num, format_context->streams[index]->time_base.den);
        printf("AVStream->start_time %ld.\n", format_context->streams[index]->start_time);
        printf("AVStream->duration %ld.\n", format_context->streams[index]->duration);

        printf("finding the proper (CODEC).\n");

        const AVCodec * local_codec = NULL;
        local_codec = avcodec_find_decoder(local_parameters->codec_id);
        if (local_codec == NULL)
        {
            printf("[ERROR][avcodec_find_decoder] unsupported codec.\n");
            continue;
        }
        if (local_parameters->codec_type == AVMEDIA_TYPE_VIDEO)
        {
            if (video_stream_index == -1)
            {
                video_stream_index = index;
                codec = local_codec;
                codec_parameters = local_parameters;
            }
            printf("Video Codec: resolution %d x %d.\n", local_parameters->width, local_parameters->height);
        }
        else if (local_parameters->codec_type == AVMEDIA_TYPE_AUDIO)
        {
            printf("Audio Codec: %d channels, sample rate %d.\n", local_parameters->channels, local_parameters->sample_rate);
        }
        printf("\tCodec %s ID %d bit_rate %ld.\n", local_codec->name, local_codec->id, local_parameters->bit_rate);
    }
    AVCodecContext * video_codec_context = NULL;
    if (video_stream_index == -1)
    {
        printf("no video stream information.\n");
        return -1;
    }
    video_codec_context = avcodec_alloc_context3(codec);
    if (video_codec_context == NULL)
    {
        printf("[ERROR][avcodec_alloc_context3] alloc context3.\n");
        return -1;
    }
    result = avcodec_parameters_to_context(video_codec_context, codec_parameters);
    if (result < 0)
    {
        printf("[ERROR][avcodec_parameters_to_context] copy codec parameter to codec context failed.\n");
        return -1;
    }
    result = avcodec_open2(video_codec_context, codec, NULL);
    if (result < 0)
    {
        printf("[ERROR][avcodec_open2] failed to copy codec params to codec context.\n");
        return -1;
    }
    AVFrame  * frame  = av_frame_alloc();
    AVPacket * packet = av_packet_alloc();
    if (frame == NULL || packet == NULL)
    {
        printf("[ERROR] av alloc failed.\n");
        return -1;
    }
    int how_many_packets_to_process = format_context->streams[video_stream_index]->r_frame_rate.num;
    while(av_read_frame(format_context, packet) >= 0)
    {
        if (packet->stream_index == video_stream_index)
        {
            printf("AVPacket->pts %ld.\n", packet->pts);
            result = decode_packet(packet, video_codec_context, frame);
            if (result < 0)
            {
                break;
            }
            if ((--how_many_packets_to_process) <= 0)
            {
                break;
            }
        }
        av_packet_unref(packet);
    }
    printf("releasing all the resources.\n");

    avformat_close_input(&format_context);
    av_packet_free(&packet);
    av_frame_free(&frame);
    avcodec_free_context(&video_codec_context);
    return 0;
}

